import json
import numpy as np
from seal import *
from app.pkg.circuits import ArithmeticCircuit, BooleanCircuit
from app.pkg.polynomial import *
from collections import namedtuple
from app.pkg.utils import NumpyEncoder

QRP = namedtuple('QRP', ['V', 'W', 'Y', 't'])
MODULUS = 2 ** 61 - 1


class SNARKVerifier:
    def __init__(self, context):
        self.context = context
        self.crs = None
        self.vrs = None
        self.h = None
        self.target_parms = None
        self._hash_cache = {}

    def setup(self, function_circuit, security_param=128):
        self.h = sample_irreducible_polynomial(security_param)
        qrp = self.circuit_to_qrp(function_circuit)
        self.crs = (self.h, qrp, self.generate_crs(qrp))
        self.vrs = (self.h, qrp, self.generate_vrs(qrp))
        self.target_parms = self.h_to_seal_params(self.context)
        return self.crs, self.vrs

    def circuit_to_qrp(self, circuit):
        if isinstance(circuit, BooleanCircuit):
            return self.dynamic_qrp_generation(circuit)
        elif isinstance(circuit, ArithmeticCircuit):
            return self.dynamic_qrp_generation(circuit)
        else:
            return super().circuit_to_qrp(circuit)

    def dynamic_qrp_generation(self, circuit):
        # 生成随机门位置点
        gate_positions = np.random.randint(0, 1000, size=len(circuit.gates))
        t = np.poly1d([1] + [-pos for pos in gate_positions])

        # 初始化V, W, Y多项式
        V = []
        W = []
        Y = []

        for idx, gate in enumerate(circuit.gates):
            # 根据门类型生成多项式
            if gate['type'] == 'add':
                v_poly = np.poly1d([1, -gate['input1']])
                w_poly = np.poly1d([1, -gate['input2']])
                y_poly = np.poly1d([1, -gate['output']])
            elif gate['type'] == 'mult':
                v_poly = np.poly1d([1, -gate['input1']])
                w_poly = np.poly1d([1, -gate['input2']])
                y_poly = np.poly1d([1, -gate['output']]) - v_poly * w_poly
            else:
                continue  # 其他门类型处理

            V.append(v_poly)
            W.append(w_poly)
            Y.append(y_poly)

        # 应用快速插值算法
        V = self.fast_interpolation_optimization(V, gate_positions)
        W = self.fast_interpolation_optimization(W, gate_positions)
        Y = self.fast_interpolation_optimization(Y, gate_positions)

        return QRP(V, W, Y, t)

    def fast_interpolation_optimization(self, polynomials, positions):
        # 分治策略预计算M(x)
        def compute_M(points):
            if len(points) == 1:
                return np.poly1d([1, -points[0]])
            mid = len(points) // 2
            left = compute_M(points[:mid])
            right = compute_M(points[mid:])
            return np.polymul(left, right)

        M = compute_M(positions)
        M_deriv = np.polyder(M)

        optimized_polys = []
        for p in polynomials:
            # 计算分子多项式
            numerator = np.polymul(p, M)
            # 计算分母预存值
            denominators = M_deriv(positions)
            # 构建插值多项式
            coeffs = []
            for i in range(len(positions)):
                coeff = p(positions[i]) / denominators[i]
                coeffs.append(coeff)
            # 应用快速插值
            optimized_poly = np.poly1d(coeffs)
            optimized_polys.append(optimized_poly)

        return optimized_polys
    def boolean_circuit_to_qrp(self, circuit):
        V = [np.poly1d([1, 0]), np.poly1d([1, 0])]
        W = [np.poly1d([1, 0]), np.poly1d([0, 1])]
        Y = [np.poly1d([0]), np.poly1d([0, 1])]
        t = np.poly1d([1, -1])
        return QRP(V, W, Y, t)

    def arithmetic_circuit_to_qrp(self, circuit):
        V = [np.poly1d([1, 0]), np.poly1d([1, 0])]
        W = [np.poly1d([1]), np.poly1d([1])]
        Y = [np.poly1d([0]), np.poly1d([0, 1])]
        t = np.poly1d([1, -1])
        return QRP(V, W, Y, t)

    def generate_crs(self, qrp):
        import secrets
        r = secrets.randbelow(2 ** 16)
        return {
            'V_evals': [v(r) % MODULUS for v in qrp.V],
            'W_evals': [w(r) % MODULUS for w in qrp.W],
            'Y_evals': [y(r) % MODULUS for y in qrp.Y],
            'r': r
        }

    def generate_vrs(self, qrp):
        import secrets
        alpha = secrets.randbelow(2 ** 16)
        return {
            't_eval': qrp.t(self.crs[2]['r']) % MODULUS,
            'alpha': alpha
        }

    def compress_ciphertext(self, ciphertext):
        if self.target_parms is None:
            self.target_parms = self.h_to_seal_params(self.context)

        target_context = SEALContext(self.target_parms)
        evaluator = Evaluator(self.context)
        target_parms_id = target_context.first_context_data().parms_id()

        try:
            evaluator.mod_switch_to_inplace(ciphertext, target_parms_id)
        except ValueError as e:
            print(f"Mod switch error: {e}")
            self.target_parms = self.h_to_seal_params(self.context)
            target_context = SEALContext(self.target_parms)
            evaluator.mod_switch_to_inplace(ciphertext, target_parms_id)

        return ciphertext

    def h_to_seal_params(self, context):
        original_degree = context.key_context_data().parms().poly_modulus_degree()
        original_coeff_mod = context.key_context_data().parms().coeff_modulus()
        parms = EncryptionParameters(scheme_type.bgv)
        parms.set_poly_modulus_degree(original_degree)
        parms.set_coeff_modulus(original_coeff_mod)
        parms.set_plain_modulus(PlainModulus.Batching(original_degree, 20))
        return parms

    def generate_proof(self, compressed_ciphertext, circuit_intermediates):
        r = self.crs[2]['r']
        qrp = self.crs[1]
        points = np.array([r, (r % 1000) + 1, (r % 1000) + 2])
        values = np.array([
            (qrp.V[0](points[0]) * qrp.W[0](points[0]) - qrp.Y[0](points[0])) % MODULUS,
            (qrp.V[0](points[1]) * qrp.W[0](points[1]) - qrp.Y[0](points[1])) % MODULUS,
            (qrp.V[0](points[2]) * qrp.W[0](points[2]) - qrp.Y[0](points[2])) % MODULUS
        ])

        polynomial = fast_interpolation(points, values)
        remainder = np.polydiv(polynomial, qrp.t)[1]

        if not np.all(remainder == 0):
            s = np.poly1d([1])
            s = np.polymul(s, polynomial)
            s = np.polydiv(s, qrp.t)[0]
            polynomial = np.polymul(s, qrp.t)

        proof = {
            'polynomial': polynomial,
            'qrp_eval': self.evaluate_qrp(polynomial, r)
        }
        return json.dumps(proof, cls=NumpyEncoder).encode()

    def evaluate_qrp(self, polynomial, r):
        qrp = self.crs[1]
        result = (qrp.V[0](r) * qrp.W[0](r) - qrp.Y[0](r)) % MODULUS
        return result

    def verify(self, proof, compressed_ciphertext):
        try:
            proof_data = json.loads(proof.decode())
            polynomial = np.poly1d(proof_data['polynomial'])
            qrp_eval = proof_data['qrp_eval']
            r = self.crs[2]['r']
            qrp = self.crs[1]

            quotient, remainder = np.polydiv(polynomial, qrp.t)
            if not np.allclose(remainder, 0, atol=1e-8):
                return False

            expected_qrp_eval = (qrp.V[0](r) * qrp.W[0](r) - qrp.Y[0](r)) % MODULUS
            return np.isclose(qrp_eval, expected_qrp_eval, atol=1e-8)

        except Exception as e:
            print(f"Verification error: {e}")
            return False

    def homomorphic_compression(self, ciphertext):
        # 获取原始参数
        original_parms = self.context.key_context_data().parms()
        original_degree = original_parms.poly_modulus_degree()
        target_degree = original_degree // 2  # 必须严格为原次数的1/2

        # 继承原系数模数的前一半素数位数（例如原为[60,60,60]，则新为[60,60]）
        original_coeff_modulus = original_parms.coeff_modulus()
        original_coeff_bits = [modulus.bit_count() for modulus in original_coeff_modulus]
        target_coeff_bits = original_coeff_bits[:len(original_coeff_bits) // 2]

        # 生成目标参数
        target_parms = EncryptionParameters(scheme_type.bgv)
        target_parms.set_poly_modulus_degree(target_degree)
        target_parms.set_coeff_modulus(CoeffModulus.Create(target_degree, target_coeff_bits))
        target_parms.set_plain_modulus(original_parms.plain_modulus())  # 明文模数必须一致
        print(f"[DEBUG] Original Degree: {original_degree}")
        print(f"[DEBUG] Target Degree: {target_degree}")
        print(f"[DEBUG] Original Coeff Bits: {original_coeff_bits}")
        print(f"[DEBUG] Target Coeff Bits: {target_coeff_bits}")
        # 创建目标上下文并链接到原参数链
        target_context = SEALContext(target_parms, True)
        evaluator = Evaluator(self.context)
        target_parms_id = target_context.first_context_data().parms_id()

        try:
            evaluator.mod_switch_to_inplace(ciphertext, target_parms_id)
        except ValueError as e:
            print(f"Mod switch error: {e}")
            # 强制重新初始化上下文
            target_context = SEALContext(target_parms)
            target_parms_id = target_context.first_context_data().parms_id()
            evaluator.mod_switch_to_inplace(ciphertext, target_parms_id)

        return ciphertext